#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File: run.py
@Time: 2023/02/21 14:41:52
@Description: Nuclei POC批量转换为Snort规则工具
@Author: nuclei2snort
@Version: 2.0.0
'''

import argparse
import binascii
import datetime
import os
import re
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import yaml

from config.config import get_config, validate_config
from config.log import logger
from tools.cn2en import translate


class NucleiToSnortConverter:
    """Nuclei POC转Snort规则转换器"""
    
    def __init__(self):
        """初始化转换器"""
        self.config = get_config()
        self.validate_environment()
        self.sid_counter = self.config['snort']['start_sid']
        
    def validate_environment(self) -> bool:
        """验证运行环境"""
        try:
            validate_config()
            logger.log('INFOR', '环境验证通过')
            return True
        except Exception as e:
            logger.log('ERROR', f'环境验证失败: {e}')
            return False
    
    @staticmethod
    def string_to_hex(input_string: str) -> str:
        """
        将字符串转换为十六进制格式
        
        Args:
            input_string (str): 待转换的字符串
            
        Returns:
            str: 十六进制格式的字符串，如 |48 65 6c 6c 6f|
        """
        try:
            input_bytes = input_string.encode('utf-8')
            hex_string = binascii.hexlify(input_bytes).decode('utf-8')
            hex_formatted = ' '.join(hex_string[i:i+2] for i in range(0, len(hex_string), 2))
            return f'|{hex_formatted}|'
        except Exception as e:
            logger.log('ERROR', f'字符串转十六进制失败: {e}')
            return input_string
    
    @staticmethod
    def check_elements_in_string(elements: List[str], target_string: str) -> Optional[str]:
        """
        检查列表中的任意元素是否在目标字符串中
        
        Args:
            elements (List[str]): 要检查的元素列表
            target_string (str): 目标字符串
            
        Returns:
            Optional[str]: 找到的第一个元素，如果没有找到则返回None
        """
        for element in elements:
            if element in target_string:
                return element
        return None
    
    def get_classtype(self, message: str) -> str:
        """
        根据消息内容确定漏洞分类类型
        
        Args:
            message (str): 漏洞描述信息
            
        Returns:
            str: Snort规则分类类型
        """
        mapping = self.config['mappings']['classtype']
        classtype = self.check_elements_in_string(list(mapping.keys()), message)
        return mapping.get(classtype, self.config['snort']['default_classtype'])
    
    def get_priority(self, severity: str) -> str:
        """
        根据漏洞严重程度获取优先级
        
        Args:
            severity (str): 漏洞严重程度
            
        Returns:
            str: 优先级数值
        """
        mapping = self.config['mappings']['priority']
        priority = self.check_elements_in_string(list(mapping.keys()), severity.lower())
        return mapping.get(priority, self.config['snort']['default_priority'])
    
    def generate_content_rules(self, contents: List[str]) -> str:
        """
        生成Snort规则的content部分
        
        Args:
            contents (List[str]): 内容列表
            
        Returns:
            str: 格式化的content规则
        """
        content_rules = []
        for content in contents:
            if content.startswith('|') and content.endswith('|'):
                # 十六进制内容
                content_rules.append(f'content:"{content}";')
            else:
                # 普通文本内容
                content_rules.append(f'content:"{content}";nocase;')
        return ''.join(content_rules)
    
    def generate_snort_rule(self, rule_data: Dict) -> str:
        """
        生成完整的Snort规则
        
        Args:
            rule_data (Dict): 规则数据字典
            
        Returns:
            str: 完整的Snort规则
        """
        template = (
            f'alert http any any -> any any ('
            f'msg:"{rule_data["msg_en"]}";'
            f'msg_cn:"{rule_data["msg_cn"]}";'
            f'flow:to_server,established;'
            f'{rule_data["content"]}'
            f'sid:{rule_data["sid"]};'
            f'priority:{rule_data["priority"]};'
            f'classtype:{rule_data["classtype"]};'
            f'cve:{rule_data["cve_id"]};'
            f'success:0;'
            f'attack_chain:intrusion;'
            f'reliability:{self.config["snort"]["reliability"]};'
            f')'
        )
        return template
    
    def save_rule_to_file(self, rule: str) -> None:
        """
        保存规则到文件
        
        Args:
            rule (str): Snort规则字符串
        """
        try:
            now = datetime.datetime.now()
            time_str = now.strftime("%Y%m%d%H")
            filename = f'{time_str}_ips.rules'
            output_file = self.config['paths']['output_path'] / filename
            
            with open(output_file, 'a', encoding='utf-8') as f:
                f.write(rule + '\n')
            
            logger.log('INFOR', f'规则已保存到: {output_file}')
        except Exception as e:
            logger.log('ERROR', f'保存规则文件失败: {e}')
    
    def extract_poc_content(self, yaml_file: Path) -> Optional[Dict]:
        """
        从YAML文件中提取POC内容
        
        Args:
            yaml_file (Path): YAML文件路径
            
        Returns:
            Optional[Dict]: 提取的内容字典，失败返回None
        """
        try:
            with open(yaml_file, 'r', encoding='utf-8') as f:
                data = yaml.safe_load(f)
            
            if not data or 'http' not in data or 'info' not in data:
                logger.log('ERROR', f'{yaml_file.name}: YAML格式不正确')
                return None
            
            # 提取基本信息
            http_data = str(data['http'])
            info = data['info']
            
            # 获取CVE ID
            cve_id = 'N/A'
            if 'classification' in info and 'cve-id' in info['classification']:
                cve_id = str(info['classification']['cve-id'])
            
            # 获取漏洞信息
            msg_en = str(info.get('name', 'Unknown Vulnerability'))
            severity = str(info.get('severity', 'medium'))
            
            # 翻译获取中文名称
            msg_cn = translate(msg_en)
            
            # 获取分类和优先级
            priority = self.get_priority(severity)
            classtype = self.get_classtype(msg_cn)
            
            # 提取标记内容
            pattern = re.compile(r'\$\$(.*?)\$\$')
            matches = re.findall(pattern, http_data)
            
            if not matches:
                logger.log('ERROR', f'{yaml_file.name}: 未找到标记内容($$...$$)')
                return None
            
            # 处理HTTP方法
            http_method = self.check_elements_in_string(
                self.config['constants']['http_methods'], 
                http_data
            )
            
            # 处理content内容
            content_list = []
            for i, match in enumerate(matches, 1):
                logger.log('INFOR', f'处理第{i}个标记: {match[:50]}...')
                
                # 检查是否包含特殊字符
                special_char = self.check_elements_in_string(
                    self.config['constants']['special_characters'], 
                    match
                )
                
                if special_char:
                    hex_content = self.string_to_hex(match)
                    content_list.append(hex_content)
                    logger.log('INFOR', f'检测到特殊字符"{special_char}"，已转换为HEX: {hex_content}')
                else:
                    content_list.append(match)
                    logger.log('INFOR', f'普通内容: {match}')
            
            return {
                'cve_id': cve_id,
                'msg_en': msg_en,
                'msg_cn': msg_cn,
                'priority': priority,
                'classtype': classtype,
                'content_list': content_list,
                'http_method': http_method
            }
            
        except yaml.YAMLError as e:
            logger.log('ERROR', f'{yaml_file.name}: YAML解析错误 - {e}')
            return None
        except FileNotFoundError:
            logger.log('ERROR', f'{yaml_file.name}: 文件不存在')
            return None
        except Exception as e:
            logger.log('ERROR', f'{yaml_file.name}: 处理失败 - {e}')
            return None
    
    def convert_single_file(self, yaml_file: Path) -> bool:
        """
        转换单个YAML文件为Snort规则
        
        Args:
            yaml_file (Path): YAML文件路径
            
        Returns:
            bool: 转换是否成功
        """
        logger.log('INFOR', f'开始处理文件: {yaml_file.name}')
        
        poc_data = self.extract_poc_content(yaml_file)
        if not poc_data:
            return False
        
        # 生成规则数据
        rule_data = {
            'msg_en': poc_data['msg_en'],
            'msg_cn': poc_data['msg_cn'],
            'classtype': poc_data['classtype'],
            'priority': poc_data['priority'],
            'sid': self.sid_counter,
            'content': self.generate_content_rules(poc_data['content_list']),
            'cve_id': poc_data['cve_id']
        }
        
        # 生成Snort规则
        snort_rule = self.generate_snort_rule(rule_data)
        
        # 保存规则
        self.save_rule_to_file(snort_rule)
        
        logger.log('ALERT', f'成功生成规则:\n{snort_rule}')
        self.sid_counter += 1
        
        return True
    
    def convert_directory(self, directory: Path, max_workers: int = 4) -> Tuple[int, int]:
        """
        批量转换目录中的所有YAML文件
        
        Args:
            directory (Path): 目标目录
            max_workers (int): 最大并发数
            
        Returns:
            Tuple[int, int]: (成功数量, 总数量)
        """
        logger.log('INFOR', f'开始批量转换目录: {directory}')
        
        # 查找所有YAML文件
        yaml_files = []
        for root, dirs, files in os.walk(directory):
            for file in files:
                if file.endswith(('.yaml', '.yml')):
                    yaml_files.append(Path(root) / file)
        
        if not yaml_files:
            logger.log('ERROR', f'目录 {directory} 中未找到YAML文件')
            return 0, 0
        
        logger.log('INFOR', f'找到 {len(yaml_files)} 个YAML文件')
        
        success_count = 0
        
        # 使用线程池进行并发处理
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # 提交所有任务
            future_to_file = {
                executor.submit(self.convert_single_file, yaml_file): yaml_file 
                for yaml_file in yaml_files
            }
            
            # 处理完成的任务
            for future in as_completed(future_to_file):
                yaml_file = future_to_file[future]
                try:
                    if future.result():
                        success_count += 1
                except Exception as e:
                    logger.log('ERROR', f'处理文件 {yaml_file.name} 时发生错误: {e}')
        
        logger.log('INFOR', f'批量转换完成: 成功 {success_count}/{len(yaml_files)}')
        return success_count, len(yaml_files)


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description='Nuclei POC到Snort规则转换工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
使用示例:
  %(prog)s --file input/pocs/2023/CVE-2023-0001.yaml    # 转换单个文件
  %(prog)s --dir input/pocs/2023/                       # 转换整个目录
  %(prog)s --dir input/cve_fail/ --workers 8            # 使用8个并发线程
        '''
    )
    
    parser.add_argument(
        '--file', '-f',
        type=str,
        help='转换单个YAML文件'
    )
    
    parser.add_argument(
        '--dir', '-d',
        type=str,
        help='转换目录中的所有YAML文件'
    )
    
    parser.add_argument(
        '--workers', '-w',
        type=int,
        default=4,
        help='并发处理线程数 (默认: 4)'
    )
    
    parser.add_argument(
        '--version', '-v',
        action='version',
        version='%(prog)s 2.0.0'
    )
    
    args = parser.parse_args()
    
    # 检查参数
    if not args.file and not args.dir:
        parser.error('必须指定 --file 或 --dir 参数')
    
    if args.file and args.dir:
        parser.error('--file 和 --dir 参数不能同时使用')
    
    # 初始化转换器
    try:
        converter = NucleiToSnortConverter()
    except Exception as e:
        logger.log('FATAL', f'初始化转换器失败: {e}')
        sys.exit(1)
    
    start_time = time.time()
    
    try:
        if args.file:
            # 转换单个文件
            file_path = Path(args.file)
            if not file_path.exists():
                logger.log('ERROR', f'文件不存在: {file_path}')
                sys.exit(1)
            
            success = converter.convert_single_file(file_path)
            if success:
                logger.log('INFOR', '单文件转换完成')
            else:
                logger.log('ERROR', '单文件转换失败')
                sys.exit(1)
        
        elif args.dir:
            # 转换目录
            dir_path = Path(args.dir)
            if not dir_path.exists() or not dir_path.is_dir():
                logger.log('ERROR', f'目录不存在: {dir_path}')
                sys.exit(1)
            
            success_count, total_count = converter.convert_directory(dir_path, args.workers)
            if success_count > 0:
                logger.log('INFOR', f'目录转换完成: {success_count}/{total_count}')
            else:
                logger.log('ERROR', '目录转换失败')
                sys.exit(1)
    
    except KeyboardInterrupt:
        logger.log('ALERT', '用户中断操作')
        sys.exit(0)
    except Exception as e:
        logger.log('FATAL', f'转换过程中发生未知错误: {e}')
        sys.exit(1)
    
    end_time = time.time()
    elapsed_time = end_time - start_time
    logger.log('INFOR', f'总耗时: {elapsed_time:.2f}秒')


if __name__ == '__main__':
    main()